import netCDF4 as nc
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.animation as animation
import seaborn as sns
from dotenv import load_dotenv
from pandas.plotting import parallel_coordinates
import importlib
import plotly.express as px
import plotly.graph_objects as go
import os
import glob
import pandas as pd
import json

import utils.db_tools as db_tools
from utils.db_tools import (
    get_db,
    filter_df,
    make_animation,
    get_data,
    metrics_grid,
    plot_grid,
    compute_metrics
)

from classify import classify_trajectories

importlib.reload(db_tools)

<module 'utils.db_tools' from '/cluster/home/vogtva/pde-solvers-cuda/analysis/utils/db_tools.py'>

model = "bruss"
run_id = "ball_big"
load_dotenv()
data_dir = os.getenv("DATA_DIR")
output_dir = os.getenv("OUT_DIR")
df = pd.read_csv(f"{output_dir}/{model}/{run_id}/classification_metrics.csv")
df_class = classify_trajectories(
    df,
    deviation_threshold=1.28
)
df = df_class.copy()
# df = df[df["filename"].apply(os.path.exists)].reset_index(drop=True)
df["op"] = df["original_point"].astype(str)

df.columns

Index(['Unnamed: 0', 'model', 'A', 'B', 'Nx', 'dx', 'Nt', 'dt', 'Du', 'Dv',
       'initial_condition', 'random_seed', 'n_snapshots', 'filename', 'run_id',
       'original_point', 'mean_deviation', 'std_deviation', 'max_dx',
       'mean_dx', 'max_dt', 'mean_dt', 'dominant_power', 'total_power',
       'max_u', 'max_v', 'category', 'op'],
      dtype='object')

# Plot the distribution of 'mean_deviation' for each category
plt.figure(figsize=(12, 8))
df_filt = df_class[df_class.max_dt > 10]
ax = sns.histplot(data=df_filt, x='max_dt', hue='category', multiple='stack', kde=False)

plt.xlabel('Mean Deviation')
plt.ylabel('Frequency')
plt.title('Distribution of Mean Deviation by Category')
plt.show()

# plt.figure(figsize=(10, 6))
# sns.scatterplot(x=df_class['A'], y=df_class['B'], hue=df_class["category"])
# plt.xlabel('A')
# plt.ylabel('B')
# plt.title('Scatter plot of A vs B for Sampling Centers')
# plt.show()

fig = px.scatter(
    df_class,
    x="A",
    y="B",
    color="category",
    title="Scatter plot of A vs B",
    labels={"A": "A", "B": "B"},
    width=800,
    height=800,
)

# Display the plot in the notebook
fig.show()

df.value_counts("category")

category
SS     3957
OSC    1530
I      1294
DSS     419
Name: count, dtype: int64

df_class["op"] = df_class["original_point"].astype(str)
for _, df1 in df_class.groupby("op"):
    original_point = df1.iloc[0]["original_point"]
    print(original_point, df1.value_counts("category").to_dict())

{'A': 0.5, 'B': 0.625, 'Du': 1, 'Dv': 11} {'SS': 60}
{'A': 0.5, 'B': 0.625, 'Du': 1, 'Dv': 18} {'SS': 60}
{'A': 0.5, 'B': 0.625, 'Du': 1, 'Dv': 4} {'SS': 60}
{'A': 0.5, 'B': 0.625, 'Du': 3, 'Dv': 12} {'SS': 60}
{'A': 0.5, 'B': 0.625, 'Du': 3, 'Dv': 33} {'SS': 60}
{'A': 0.5, 'B': 0.625, 'Du': 3, 'Dv': 54} {'SS': 60}
{'A': 0.5, 'B': 1.0, 'Du': 1, 'Dv': 11} {'SS': 60}
{'A': 0.5, 'B': 1.0, 'Du': 1, 'Dv': 18} {'SS': 60}
{'A': 0.5, 'B': 1.0, 'Du': 1, 'Dv': 4} {'SS': 60}
{'A': 0.5, 'B': 1.0, 'Du': 3, 'Dv': 12} {'SS': 60}
{'A': 0.5, 'B': 1.0, 'Du': 3, 'Dv': 33} {'SS': 60}
{'A': 0.5, 'B': 1.0, 'Du': 3, 'Dv': 54} {'SS': 60}
{'A': 0.5, 'B': 1.5, 'Du': 1, 'Dv': 11} {'SS': 59, 'OSC': 1}
{'A': 0.5, 'B': 1.5, 'Du': 1, 'Dv': 18} {'SS': 47, 'I': 11, 'OSC': 2}
{'A': 0.5, 'B': 1.5, 'Du': 1, 'Dv': 4} {'SS': 60}
{'A': 0.5, 'B': 1.5, 'Du': 3, 'Dv': 12} {'SS': 60}
{'A': 0.5, 'B': 1.5, 'Du': 3, 'Dv': 33} {'SS': 60}
{'A': 0.5, 'B': 1.5, 'Du': 3, 'Dv': 54} {'SS': 54, 'I': 5, 'OSC': 1}
{'A': 0.5, 'B': 2.0, 'Du': 1, 'Dv': 11} {'OSC': 43, 'SS': 17}
{'A': 0.5, 'B': 2.0, 'Du': 1, 'Dv': 18} {'OSC': 41, 'SS': 10, 'I': 9}
{'A': 0.5, 'B': 2.0, 'Du': 1, 'Dv': 4} {'OSC': 48, 'SS': 12}
{'A': 0.5, 'B': 2.0, 'Du': 3, 'Dv': 12} {'OSC': 49, 'SS': 11}
{'A': 0.5, 'B': 2.0, 'Du': 3, 'Dv': 33} {'OSC': 49, 'SS': 11}
{'A': 0.5, 'B': 2.0, 'Du': 3, 'Dv': 54} {'OSC': 40, 'SS': 19, 'I': 1}
{'A': 1.0, 'B': 1.25, 'Du': 1, 'Dv': 11} {'SS': 60}
{'A': 1.0, 'B': 1.25, 'Du': 1, 'Dv': 18} {'SS': 60}
{'A': 1.0, 'B': 1.25, 'Du': 1, 'Dv': 4} {'SS': 60}
{'A': 1.0, 'B': 1.25, 'Du': 3, 'Dv': 12} {'SS': 60}
{'A': 1.0, 'B': 1.25, 'Du': 3, 'Dv': 33} {'SS': 60}
{'A': 1.0, 'B': 1.25, 'Du': 3, 'Dv': 54} {'SS': 60}
{'A': 1.0, 'B': 2.0, 'Du': 1, 'Dv': 11} {'SS': 59, 'I': 1}
{'A': 1.0, 'B': 2.0, 'Du': 1, 'Dv': 18} {'SS': 40, 'I': 19, 'DSS': 1}
{'A': 1.0, 'B': 2.0, 'Du': 1, 'Dv': 4} {'SS': 60}
{'A': 1.0, 'B': 2.0, 'Du': 3, 'Dv': 12} {'SS': 60}
{'A': 1.0, 'B': 2.0, 'Du': 3, 'Dv': 33} {'SS': 58, 'OSC': 2}
{'A': 1.0, 'B': 2.0, 'Du': 3, 'Dv': 54} {'SS': 41, 'I': 19}
{'A': 1.0, 'B': 3.0, 'Du': 1, 'Dv': 11} {'OSC': 32, 'I': 24, 'DSS': 4}
{'A': 1.0, 'B': 3.0, 'Du': 1, 'Dv': 18} {'OSC': 29, 'I': 24, 'DSS': 6, 'SS': 1}
{'A': 1.0, 'B': 3.0, 'Du': 1, 'Dv': 4} {'OSC': 49, 'SS': 10, 'I': 1}
{'A': 1.0, 'B': 3.0, 'Du': 3, 'Dv': 12} {'OSC': 56, 'SS': 3, 'I': 1}
{'A': 1.0, 'B': 3.0, 'Du': 3, 'Dv': 33} {'OSC': 43, 'DSS': 9, 'I': 8}
{'A': 1.0, 'B': 3.0, 'Du': 3, 'Dv': 54} {'I': 30, 'OSC': 29, 'SS': 1}
{'A': 1.0, 'B': 4.0, 'Du': 1, 'Dv': 11} {'OSC': 30, 'I': 20, 'DSS': 10}
{'A': 1.0, 'B': 4.0, 'Du': 1, 'Dv': 18} {'OSC': 30, 'DSS': 22, 'I': 8}
{'A': 1.0, 'B': 4.0, 'Du': 1, 'Dv': 4} {'OSC': 46, 'I': 14}
{'A': 1.0, 'B': 4.0, 'Du': 3, 'Dv': 12} {'OSC': 57, 'I': 3}
{'A': 1.0, 'B': 4.0, 'Du': 3, 'Dv': 33} {'OSC': 34, 'I': 26}
{'A': 1.0, 'B': 4.0, 'Du': 3, 'Dv': 54} {'OSC': 30, 'I': 28, 'DSS': 2}
{'A': 1.5, 'B': 1.875, 'Du': 1, 'Dv': 11} {'SS': 60}
{'A': 1.5, 'B': 1.875, 'Du': 1, 'Dv': 18} {'SS': 60}
{'A': 1.5, 'B': 1.875, 'Du': 1, 'Dv': 4} {'SS': 60}
{'A': 1.5, 'B': 1.875, 'Du': 3, 'Dv': 12} {'SS': 60}
{'A': 1.5, 'B': 1.875, 'Du': 3, 'Dv': 33} {'SS': 60}
{'A': 1.5, 'B': 1.875, 'Du': 3, 'Dv': 54} {'SS': 60}
{'A': 1.5, 'B': 3.0, 'Du': 1, 'Dv': 11} {'SS': 39, 'I': 21}
{'A': 1.5, 'B': 3.0, 'Du': 1, 'Dv': 18} {'I': 30, 'SS': 30}
{'A': 1.5, 'B': 3.0, 'Du': 1, 'Dv': 4} {'SS': 58, 'OSC': 2}
{'A': 1.5, 'B': 3.0, 'Du': 3, 'Dv': 12} {'SS': 59, 'OSC': 1}
{'A': 1.5, 'B': 3.0, 'Du': 3, 'Dv': 33} {'SS': 41, 'I': 19}
{'A': 1.5, 'B': 3.0, 'Du': 3, 'Dv': 54} {'I': 30, 'SS': 28, 'OSC': 2}
{'A': 1.5, 'B': 4.5, 'Du': 1, 'Dv': 11} {'OSC': 30, 'I': 27, 'DSS': 3}
{'A': 1.5, 'B': 4.5, 'Du': 1, 'Dv': 18} {'OSC': 30, 'I': 22, 'DSS': 8}
{'A': 1.5, 'B': 4.5, 'Du': 1, 'Dv': 4} {'OSC': 43, 'I': 10, 'SS': 6, 'DSS': 1}
{'A': 1.5, 'B': 4.5, 'Du': 3, 'Dv': 12} {'OSC': 45, 'I': 9, 'SS': 6}
{'A': 1.5, 'B': 4.5, 'Du': 3, 'Dv': 33} {'I': 30, 'OSC': 30}
{'A': 1.5, 'B': 4.5, 'Du': 3, 'Dv': 54} {'I': 30, 'OSC': 29, 'SS': 1}
{'A': 1.5, 'B': 6.0, 'Du': 1, 'Dv': 11} {'OSC': 30, 'DSS': 25, 'I': 5}
{'A': 1.5, 'B': 6.0, 'Du': 1, 'Dv': 18} {'OSC': 30, 'DSS': 26, 'I': 4}
{'A': 1.5, 'B': 6.0, 'Du': 1, 'Dv': 4} {'OSC': 44, 'I': 11, 'DSS': 5}
{'A': 1.5, 'B': 6.0, 'Du': 3, 'Dv': 12} {'OSC': 45, 'I': 15}
{'A': 1.5, 'B': 6.0, 'Du': 3, 'Dv': 33} {'I': 30, 'OSC': 30}
{'A': 1.5, 'B': 6.0, 'Du': 3, 'Dv': 54} {'I': 30, 'OSC': 30}
{'A': 2.0, 'B': 2.5, 'Du': 1, 'Dv': 11} {'SS': 60}
{'A': 2.0, 'B': 2.5, 'Du': 1, 'Dv': 18} {'SS': 55, 'I': 5}
{'A': 2.0, 'B': 2.5, 'Du': 1, 'Dv': 4} {'SS': 60}
{'A': 2.0, 'B': 2.5, 'Du': 3, 'Dv': 12} {'SS': 60}
{'A': 2.0, 'B': 2.5, 'Du': 3, 'Dv': 33} {'SS': 60}
{'A': 2.0, 'B': 2.5, 'Du': 3, 'Dv': 54} {'SS': 57, 'I': 3}
{'A': 2.0, 'B': 4.0, 'Du': 1, 'Dv': 11} {'I': 30, 'SS': 30}
{'A': 2.0, 'B': 4.0, 'Du': 1, 'Dv': 18} {'I': 30, 'SS': 30}
{'A': 2.0, 'B': 4.0, 'Du': 1, 'Dv': 4} {'SS': 60}
{'A': 2.0, 'B': 4.0, 'Du': 3, 'Dv': 12} {'SS': 60}
{'A': 2.0, 'B': 4.0, 'Du': 3, 'Dv': 33} {'I': 30, 'SS': 30}
{'A': 2.0, 'B': 4.0, 'Du': 3, 'Dv': 54} {'I': 30, 'SS': 30}
{'A': 2.0, 'B': 6.0, 'Du': 1, 'Dv': 11} {'OSC': 28, 'DSS': 21, 'I': 9, 'SS': 2}
{'A': 2.0, 'B': 6.0, 'Du': 1, 'Dv': 18} {'OSC': 30, 'DSS': 21, 'I': 9}
{'A': 2.0, 'B': 6.0, 'Du': 1, 'Dv': 4} {'OSC': 33, 'I': 22, 'DSS': 3, 'SS': 2}
{'A': 2.0, 'B': 6.0, 'Du': 3, 'Dv': 12} {'OSC': 33, 'I': 24, 'SS': 3}
{'A': 2.0, 'B': 6.0, 'Du': 3, 'Dv': 33} {'I': 30, 'OSC': 29, 'SS': 1}
{'A': 2.0, 'B': 6.0, 'Du': 3, 'Dv': 54} {'I': 30, 'OSC': 18, 'SS': 12}
{'A': 2.0, 'B': 8.0, 'Du': 1, 'Dv': 11} {'OSC': 30, 'DSS': 23, 'I': 7}
{'A': 2.0, 'B': 8.0, 'Du': 1, 'Dv': 18} {'OSC': 30, 'DSS': 29, 'I': 1}
{'A': 2.0, 'B': 8.0, 'Du': 1, 'Dv': 4} {'OSC': 32, 'DSS': 22, 'I': 6}
{'A': 2.0, 'B': 8.0, 'Du': 3, 'Dv': 12} {'OSC': 45, 'I': 15}
{'A': 2.0, 'B': 8.0, 'Du': 3, 'Dv': 33} {'OSC': 30, 'I': 28, 'DSS': 2}
{'A': 2.0, 'B': 8.0, 'Du': 3, 'Dv': 54} {'I': 30, 'OSC': 30}
{'A': 5.0, 'B': 10.0, 'Du': 1, 'Dv': 11} {'SS': 30, 'I': 28, 'DSS': 2}
{'A': 5.0, 'B': 10.0, 'Du': 1, 'Dv': 18} {'SS': 30, 'I': 22, 'DSS': 8}
{'A': 5.0, 'B': 10.0, 'Du': 1, 'Dv': 4} {'SS': 60}
{'A': 5.0, 'B': 10.0, 'Du': 3, 'Dv': 12} {'SS': 60}
{'A': 5.0, 'B': 10.0, 'Du': 3, 'Dv': 33} {'I': 30, 'SS': 30}
{'A': 5.0, 'B': 10.0, 'Du': 3, 'Dv': 54} {'I': 30, 'SS': 30}
{'A': 5.0, 'B': 15.0, 'Du': 1, 'Dv': 11} {'SS': 30, 'DSS': 21, 'I': 9}
{'A': 5.0, 'B': 15.0, 'Du': 1, 'Dv': 18} {'SS': 30, 'DSS': 26, 'I': 4}
{'A': 5.0, 'B': 15.0, 'Du': 1, 'Dv': 4} {'SS': 33, 'I': 14, 'DSS': 13}
{'A': 5.0, 'B': 15.0, 'Du': 3, 'Dv': 12} {'SS': 31, 'I': 29}
{'A': 5.0, 'B': 15.0, 'Du': 3, 'Dv': 33} {'SS': 30, 'I': 25, 'DSS': 5}
{'A': 5.0, 'B': 15.0, 'Du': 3, 'Dv': 54} {'SS': 30, 'I': 22, 'DSS': 8}
{'A': 5.0, 'B': 20.0, 'Du': 1, 'Dv': 11} {'SS': 30, 'DSS': 29, 'I': 1}
{'A': 5.0, 'B': 20.0, 'Du': 1, 'Dv': 18} {'SS': 30, 'I': 21, 'DSS': 9}
{'A': 5.0, 'B': 20.0, 'Du': 1, 'Dv': 4} {'SS': 30, 'DSS': 28, 'I': 2}
{'A': 5.0, 'B': 20.0, 'Du': 3, 'Dv': 12} {'SS': 30, 'I': 26, 'DSS': 4}
{'A': 5.0, 'B': 20.0, 'Du': 3, 'Dv': 33} {'SS': 30, 'DSS': 21, 'I': 9}
{'A': 5.0, 'B': 20.0, 'Du': 3, 'Dv': 54} {'I': 30, 'SS': 30}
{'A': 5.0, 'B': 6.25, 'Du': 1, 'Dv': 11} {'SS': 52, 'I': 8}
{'A': 5.0, 'B': 6.25, 'Du': 1, 'Dv': 18} {'I': 30, 'SS': 30}
{'A': 5.0, 'B': 6.25, 'Du': 1, 'Dv': 4} {'SS': 60}
{'A': 5.0, 'B': 6.25, 'Du': 3, 'Dv': 12} {'SS': 60}
{'A': 5.0, 'B': 6.25, 'Du': 3, 'Dv': 33} {'SS': 48, 'I': 10, 'DSS': 2}
{'A': 5.0, 'B': 6.25, 'Du': 3, 'Dv': 54} {'I': 30, 'SS': 30}

def plot_ball_behavior(df, metric="dev"):
    t = np.linspace(0, 100, 100)
    title = ""

    all_metrics = []
    for _, row in df.iterrows():
        d = get_data(row)
        metrics = compute_metrics(row, start_frame=0)
        if metric == "dev":
            title = "Deviation"
            values = metrics[0]
        elif metric == "dt":
            title = "Time Derivative"
            values = metrics[1]
        elif metric == "dx":
            title = "Spatial Derivative"
            values = metrics[2]
        all_metrics.append(values)
    
    # Convert to numpy array for easier computation
    all_metrics = np.array(all_metrics)

    # Compute mean and std
    avg_metric = np.mean(all_metrics, axis=0)
    std_metric = np.std(all_metrics, axis=0)

    # Create figure
    fig = go.Figure()

    # Add shaded area for standard deviation
    fig.add_trace(
        go.Scatter(
            x=np.concatenate([t, t[::-1]]),
            y=np.concatenate(
                [avg_metric + std_metric, (avg_metric - std_metric)[::-1]]
            ),
            fill="toself",
            fillcolor="rgba(0,100,80,0.2)",
            line=dict(color="rgba(255,255,255,0)"),
            showlegend=False,
        )
    )

    # Add mean line
    fig.add_trace(
        go.Scatter(
            x=t,
            y=avg_metric,
            mode="lines",
            name=title,
            hovertemplate="Index: %{x}<br>Deviation: %{y:.2f}<extra></extra>",
        )
    )

    # Update layout
    fig.update_layout(
        title="Deviation Metrics",
        xaxis_title="Time Step/Index",
        yaxis_title="Deviation Value",
        hovermode="x unified",
        showlegend=True,
        template="plotly_white",
    )

    fig.show()

for _, df1 in df_class.groupby("op"):
    print(df1.iloc[0]["op"])
    plot_ball_behavior(df1)

{'A': 0.5, 'B': 0.625, 'Du': 1, 'Dv': 11}

{'A': 0.5, 'B': 0.625, 'Du': 1, 'Dv': 18}

{'A': 0.5, 'B': 0.625, 'Du': 1, 'Dv': 4}

{'A': 0.5, 'B': 0.625, 'Du': 3, 'Dv': 12}

{'A': 0.5, 'B': 0.625, 'Du': 3, 'Dv': 33}

{'A': 0.5, 'B': 0.625, 'Du': 3, 'Dv': 54}

{'A': 0.5, 'B': 1.0, 'Du': 1, 'Dv': 11}

{'A': 0.5, 'B': 1.0, 'Du': 1, 'Dv': 18}

{'A': 0.5, 'B': 1.0, 'Du': 1, 'Dv': 4}

{'A': 0.5, 'B': 1.0, 'Du': 3, 'Dv': 12}